import cv2
import os
import numpy as np
from pathlib import Path
import shutil
import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter import ttk

class ImageMatcher:
    VALID_EXTENSIONS = {'.jpg', '.jpeg', '.png'}
    
    def __init__(self, logos_dir, target_dir, output_base_dir):
        self.invoices_dir = Path(logos_dir)
        self.target_dir = Path(target_dir)
        self.output_base_dir = Path(output_base_dir)
        self.invoices = {}
        self.load_invoices()
        
    def load_invoices(self):
        """加载所有发票模板图片"""
        for invoice_path in self.invoices_dir.glob('*'):
            if invoice_path.suffix.lower() in self.VALID_EXTENSIONS:
                invoice_img = cv2.imdecode(np.fromfile(str(invoice_path), dtype=np.uint8), cv2.IMREAD_COLOR)
                if invoice_img is not None:
                    self.invoices[invoice_path.stem] = invoice_img
        
    def match_image(self, img, invoice):
        """特征匹配函数"""
        try:
            # 转换为灰度图像可能会提高匹配效果
            img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            invoice_gray = cv2.cvtColor(invoice, cv2.COLOR_BGR2GRAY)
            
            sift = cv2.SIFT_create()
            kp1, des1 = sift.detectAndCompute(invoice_gray, None)
            kp2, des2 = sift.detectAndCompute(img_gray, None)
            
            if des1 is None or des2 is None or len(des1) < 2 or len(des2) < 2:
                return 0
            
            flann = cv2.FlannBasedMatcher(
                dict(algorithm=1, trees=5),
                dict(checks=50)
            )
            
            matches = flann.knnMatch(des1, des2, k=2)
            good_matches = [m for m, n in matches if m.distance < 0.7 * n.distance]
            return len(good_matches)
        except Exception as e:
            print(f"匹配过程出错: {str(e)}")
            return 0
        
    def process_images(self):
        """处理图片函数"""
        if not self.invoices:
            return
            
        # 创建输出目录
        for invoice_name in self.invoices:
            (self.output_base_dir / invoice_name).mkdir(parents=True, exist_ok=True)
            
        # 处理图片
        for img_path in self.target_dir.glob('*'):
            if img_path.suffix.lower() not in self.VALID_EXTENSIONS:
                continue
                
            img = cv2.imdecode(np.fromfile(str(img_path), dtype=np.uint8), cv2.IMREAD_COLOR)
            if img is None:
                continue
                
            # 寻找最佳匹配
            best_invoice = max(
                ((name, self.match_image(img, invoice))
                 for name, invoice in self.invoices.items()),
                key=lambda x: x[1],
                default=(None, 0)
            )
            
            # 如果匹配度超过阈值，复制文件
            if best_invoice[1] > 5:
                output_path = self.output_base_dir / best_invoice[0] / img_path.name
                shutil.copy2(img_path, output_path)

class ImageRecognitionGUI:
    def __init__(self):
        self.window = tk.Tk()
        self.window.title("图像识别分类工具")
        self.window.geometry("600x400")
        
        # 创建路径变量
        self.invoices_path = tk.StringVar()
        self.target_path = tk.StringVar()
        self.output_path = tk.StringVar()
        
        self.create_widgets()
        
    def create_widgets(self):
        # 创建路径选择框架
        paths_frame = ttk.LabelFrame(self.window, text="路径设置", padding="10")
        paths_frame.pack(fill="x", padx=10, pady=5)
        
        # 发票模板路径
        ttk.Label(paths_frame, text="发票模板目录:").grid(row=0, column=0, sticky="w")
        ttk.Entry(paths_frame, textvariable=self.invoices_path, width=50).grid(row=0, column=1, padx=5)
        ttk.Button(paths_frame, text="浏览", command=lambda: self.browse_directory(self.invoices_path)).grid(row=0, column=2)
        
        # 待识别图片路径
        ttk.Label(paths_frame, text="待识别图片目录:").grid(row=1, column=0, sticky="w", pady=5)
        ttk.Entry(paths_frame, textvariable=self.target_path, width=50).grid(row=1, column=1, padx=5, pady=5)
        ttk.Button(paths_frame, text="浏览", command=lambda: self.browse_directory(self.target_path)).grid(row=1, column=2, pady=5)
        
        # 输出目录路径
        ttk.Label(paths_frame, text="输出目录:").grid(row=2, column=0, sticky="w")
        ttk.Entry(paths_frame, textvariable=self.output_path, width=50).grid(row=2, column=1, padx=5)
        ttk.Button(paths_frame, text="浏览", command=lambda: self.browse_directory(self.output_path)).grid(row=2, column=2)
        
        # 开始处理按钮
        ttk.Button(self.window, text="开始处理", command=self.process_images).pack(pady=20)
        
        # 进度条
        self.progress = ttk.Progressbar(self.window, mode='indeterminate')
        self.progress.pack(fill="x", padx=10, pady=5)
        
    def browse_directory(self, path_var):
        directory = filedialog.askdirectory()
        if directory:
            path_var.set(directory)
            
    def process_images(self):
        # 验证路径
        if not all([self.invoices_path.get(), self.target_path.get(), self.output_path.get()]):
            messagebox.showerror("错误", "请选择所有必要的目录路径")
            return
            
        try:
            self.progress.start()
            matcher = ImageMatcher(
                self.invoices_path.get(),
                self.target_path.get(),
                self.output_path.get()
            )
            matcher.process_images()
            self.progress.stop()
            messagebox.showinfo("成功", "图片处理完成")
        except Exception as e:
            self.progress.stop()
            messagebox.showerror("错误", f"处理过程出错: {str(e)}")
            
    def run(self):
        self.window.mainloop()

def main():
    gui = ImageRecognitionGUI()
    gui.run()

if __name__ == "__main__":
    main()
